# predict.py
import pandas as pd
import numpy as np
import argparse
import os
import joblib
import pickle
import json
import tensorflow as tf

# 尝试从本地文件导入特征提取脚本的功能
try:
    import feature_extraction 
    if not hasattr(feature_extraction, 'AMINO_ACIDS') or not feature_extraction.AMINO_ACIDS:
        print("补丁: feature_extraction.AMINO_ACIDS 未定义或为空，使用默认值。")
        feature_extraction.AMINO_ACIDS = ['A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K', 'M', 'F', 'P', 'S', 'T', 'W', 'Y', 'V', 'X']
    if not hasattr(feature_extraction, 'one_hot_encode_first3aa'):
        raise ImportError("feature_extraction 模块缺少 one_hot_encode_first3aa 函数。")

except ImportError:
    print("警告: feature_extraction.py 未在同一目录或PYTHONPATH中找到。预测可能失败。")
    # 定义一个临时的 AMINO_ACIDS 列表以允许脚本至少尝试运行，但这是不推荐的。
    # 实际使用中应确保 feature_extraction.py 可用。
    class MockFeatureExtraction:
        AMINO_ACIDS = ['A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K', 'M', 'F', 'P', 'S', 'T', 'W', 'Y', 'V', 'X']
        def one_hot_encode_first3aa(self, df, amino_acids_list_param, fitted_encoder=None):
            # 这是一个非常简化的占位符，实际应用中需要真实的编码逻辑
            print("警告: 正在使用模拟的 one_hot_encode_first3aa。结果可能不准确。")
            if 'First3AA' not in df.columns:
                raise ValueError("Mock one_hot_encode: 'First3AA' column missing.")
            
            # 假设 fitted_encoder 是一个有效的 OneHotEncoder
            if fitted_encoder is None:
                raise ValueError("Mock one_hot_encode: fitted_encoder is required for prediction.")

            # 模拟分解和编码
            aa_cols_temp = []
            for i in range(3):
                col_name = f'AA{i+1}_pred_mock'
                df[col_name] = df['First3AA'].apply(lambda x: x[i] if len(x) == 3 and x[i] in amino_acids_list_param else 'X')
                aa_cols_temp.append(col_name)
            
            try:
                encoded_aa_mock = fitted_encoder.transform(df[aa_cols_temp])
                encoded_feature_names_mock = []
                for i in range(3):
                    for aa in amino_acids_list_param: # 使用传入的氨基酸列表
                        encoded_feature_names_mock.append(f'AA{i+1}_{aa}') # 与 feature_extraction.py 中的命名方式一致
                
                # 确保生成的列名数量与编码器的特征数量一致
                if encoded_aa_mock.shape[1] != len(encoded_feature_names_mock):
                    # 如果不一致，可能是因为 fitted_encoder 的 categories 与 amino_acids_list_param 不匹配
                    # 或者 fitted_encoder 期望的输入列与 aa_cols_temp 不匹配
                    # 对于预测，我们应该依赖 fitted_encoder.get_feature_names_out() (如果可用)
                    # 或确保 categories 与 amino_acids_list_param 完全一致
                    if hasattr(fitted_encoder, 'get_feature_names_out'):
                         encoded_feature_names_mock = fitted_encoder.get_feature_names_out([f'AA{j+1}' for j in range(3)])
                    else: # Fallback if get_feature_names_out is not available or logic is simpler
                         pass # Keep the generated names, but there might be a mismatch if not careful


                encoded_df_mock = pd.DataFrame(encoded_aa_mock, columns=encoded_feature_names_mock[:encoded_aa_mock.shape[1]], index=df.index)

                return encoded_df_mock, fitted_encoder
            except Exception as e:
                raise ValueError(f"Mock one_hot_encode transform error: {e}. Ensure encoder was trained on similar AA columns.")


    feature_extraction = MockFeatureExtraction()


def load_model(model_path, model_type):
    """加载训练好的模型。"""
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"模型文件未找到: {model_path}")
    
    if model_type in ['svm', 'rf']:
        try:
            model = joblib.load(model_path)
        except Exception as e:
            raise IOError(f"加载 scikit-learn 模型 ({model_type}) 失败: {e}")
    elif model_type in ['cnn', 'lstm', 'attention', 'bert_placeholder']:
        try:
            # 不需要 compile=False，因为我们只是进行预测，但如果模型包含自定义对象，可能需要 custom_objects
            model = tf.keras.models.load_model(model_path, compile=True)
        except Exception as e:
            raise IOError(f"加载 TensorFlow 模型 ({model_type}) 失败: {e}")
    else:
        raise ValueError(f"不支持的模型类型: {model_type}")
    print(f"{model_type.upper()} 模型从 {model_path} 加载成功。")
    return model

def preprocess_for_prediction(df, aa_encoder_path, scaler_path, feature_columns_path):
    """
    为预测准备输入数据。
    1. 加载编码器、scaler和特征列。
    2. 对'First3AA'进行独热编码。
    3. 对齐特征列。
    4. 标准化特征。
    """
    if not os.path.exists(aa_encoder_path):
        raise FileNotFoundError(f"氨基酸编码器文件未找到: {aa_encoder_path}")
    if not os.path.exists(scaler_path):
        raise FileNotFoundError(f"Scaler文件未找到: {scaler_path}")
    if not os.path.exists(feature_columns_path):
        raise FileNotFoundError(f"特征列文件未找到: {feature_columns_path}")

    try:
        with open(aa_encoder_path, 'rb') as f:
            aa_encoder = pickle.load(f)
    except Exception as e:
        raise IOError(f"加载氨基酸编码器失败: {e}")
    
    try:
        scaler = joblib.load(scaler_path)
    except Exception as e:
        raise IOError(f"加载Scaler失败: {e}")
        
    try:
        with open(feature_columns_path, 'r') as f:
            expected_feature_columns = json.load(f)
    except Exception as e:
        raise IOError(f"加载特征列文件失败: {e}")

    print("预处理器加载成功。")

    # 1. 对 First3AA 进行独热编码
    # 确保 feature_extraction.one_hot_encode_first3aa 可以处理已拟合的编码器
    # 它应该返回编码后的DataFrame和编码器本身 (尽管我们已经加载了它)
    # feature_extraction.py 中的 one_hot_encode_first3aa 需要能够接受 fitted_encoder
    
    # 检查 feature_extraction 模块是否正常加载并且函数存在
    if not hasattr(feature_extraction, 'one_hot_encode_first3aa') or not callable(feature_extraction.one_hot_encode_first3aa):
        raise RuntimeError("feature_extraction.one_hot_encode_first3aa 函数不可用。请确保 feature_extraction.py 正确。")

    df_copy = df.copy()
    # one_hot_encode_first3aa 应该使用 feature_extraction.AMINO_ACIDS
    # 并且当传入 fitted_encoder 时，它应该使用该编码器进行转换
    encoded_aa_df, _ = feature_extraction.one_hot_encode_first3aa(df_copy, 
                                                                  feature_extraction.AMINO_ACIDS, 
                                                                  fitted_encoder=aa_encoder)
    print(f"First3AA 独热编码完成，生成特征数: {encoded_aa_df.shape[1]}")


    # 2. 准备其他特征 (如果原始数据中有，并且在训练时也使用了它们)
    # 假设除了 'First3AA' 和一些元数据列外，没有其他原始特征被使用。
    # 如果有，需要在这里包含它们。
    # 例如: X_other_features = df_copy.drop(columns=['Peptide_Sequence', 'First3AA', 'log2FoldChange', 'No_DPP4', 'DPP4_Digested'], errors='ignore')
    # 目前我们的特征主要来自 First3AA 的编码。
    
    # 合并特征 - 假设 encoded_aa_df 是唯一的特征来源 (除了训练时可能存在的其他非序列特征)
    # 如果训练时有其他特征，这里的合并逻辑需要调整
    X_processed = encoded_aa_df 

    # 3. 对齐特征列 (非常重要!)
    # 使用训练时保存的列顺序和集合，对当前处理的数据进行重排和填充缺失列 (用0)
    # 确保 X_processed 中的列名与 expected_feature_columns 中的名称格式一致。
    # feature_extraction.one_hot_encode_first3aa 生成的列名应与训练时一致。
    
    # 确保 X_processed 的索引与原始 df_copy 一致，以便后续合并预测结果
    X_processed = X_processed.reindex(columns=expected_feature_columns, fill_value=0)
    print(f"特征已对齐到训练时的 {len(expected_feature_columns)} 列。")
    if X_processed.shape[1] != len(expected_feature_columns):
        print(f"警告: 对齐后的特征数量 ({X_processed.shape[1]}) 与预期 ({len(expected_feature_columns)}) 不符。")
        print(f"X_processed columns: {X_processed.columns.tolist()[:10]}...") # 打印部分列名以供调试
        print(f"Expected columns: {expected_feature_columns[:10]}...")


    # 4. 标准化特征
    X_scaled = scaler.transform(X_processed)
    print("特征标准化完成。")
    
    return X_scaled, X_processed.columns # 返回列名用于调试或验证

def predict(model, X_scaled):
    """使用加载的模型进行预测。"""
    if hasattr(model, 'predict_proba'): # scikit-learn models
        y_pred_proba = model.predict_proba(X_scaled)[:, 1] # Probability of positive class
        y_pred_class = (y_pred_proba > 0.45).astype(int)
    elif hasattr(model, 'predict'): # TensorFlow models
        y_pred_proba = model.predict(X_scaled).ravel() # Flatten if output shape is (n, 1)
        y_pred_class = (y_pred_proba > 0.45).astype(int)
    else:
        raise TypeError("模型没有 'predict' 或 'predict_proba' 方法。")
    
    return y_pred_class, y_pred_proba

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="使用训练好的模型预测肽的DPP4酶切耐受性。")
    parser.add_argument("--model_path", type=str, required=True, help="已训练模型的路径 (.h5 或 .joblib)。")
    parser.add_argument("--model_type", type=str, required=True, 
                        choices=['svm', 'rf', 'cnn', 'lstm', 'attention', 'bert_placeholder'], 
                        help="模型的类型。")
    parser.add_argument("--input_csv", type=str, required=True, help="包含待预测肽序列的输入CSV文件。")
    parser.add_argument("--output_csv", type=str, required=True, help="保存预测结果的输出CSV文件路径。")
    
    # 路径参数应该指向训练时特定模型运行的输出目录中的文件
    parser.add_argument("--model_run_dir", type=str, required=True, 
                        help="包含scaler, encoder和feature_columns的已训练模型的运行目录。")

    args = parser.parse_args()

    try:
        # 从 model_run_dir 构建 scaler, encoder, feature_columns 的路径
        # 这些文件名是在 train_model.py 中定义的
        scaler_path = os.path.join(args.model_run_dir, f"{args.model_type}_scaler.joblib")
        aa_encoder_path = os.path.join(args.model_run_dir, f"{args.model_type}_aa_encoder.pkl")
        feature_columns_path = os.path.join(args.model_run_dir, f"{args.model_type}_feature_columns.json")

        print(f"开始预测流程...")
        print(f"  模型路径: {args.model_path}")
        print(f"  模型类型: {args.model_type}")
        print(f"  输入CSV: {args.input_csv}")
        print(f"  输出CSV: {args.output_csv}")
        print(f"  模型运行目录 (用于辅助文件): {args.model_run_dir}")
        print(f"  预期Scaler路径: {scaler_path}")
        print(f"  预期Encoder路径: {aa_encoder_path}")
        print(f"  预期特征列路径: {feature_columns_path}")


        # 1. 加载原始数据
        if not os.path.exists(args.input_csv):
            raise FileNotFoundError(f"输入CSV文件未找到: {args.input_csv}")
        input_df_original = pd.read_csv(args.input_csv)
        
        # 2. 预处理数据以进行预测
        # 确保输入数据包含 'First3AA' 列
        if 'First3AA' not in input_df_original.columns:
            # 如果 DPP-IV_test.csv 有 Peptide_Sequence 列但没有 First3AA，可以尝试从中提取
            if 'Peptide_Sequence' in input_df_original.columns:
                print("输入数据中缺少 'First3AA' 列，尝试从 'Peptide_Sequence' 的前三个字符创建。")
                input_df_original['First3AA'] = input_df_original['Peptide_Sequence'].astype(str).str[:3]
                # 确保长度为3，不足则填充（例如用'X'），但这应在特征提取中处理
                input_df_original['First3AA'] = input_df_original['First3AA'].apply(lambda x: x.ljust(3, 'X') if len(x) < 3 else x)

            else:
                raise ValueError("输入CSV文件中必须包含 'First3AA' 列或 'Peptide_Sequence' 列。")

        X_predict_scaled, processed_cols = preprocess_for_prediction(input_df_original, 
                                                                     aa_encoder_path, 
                                                                     scaler_path, 
                                                                     feature_columns_path)
        print(f"数据预处理完成，用于预测的特征形状: {X_predict_scaled.shape}")
        if X_predict_scaled.shape[0] == 0:
            raise ValueError("预处理后没有数据可供预测。")

        # 3. 加载模型
        model = load_model(args.model_path, args.model_type)

        # 4. 进行预测
        predicted_classes, predicted_probabilities = predict(model, X_predict_scaled)
        print("预测完成。")

        # 5. 将预测结果添加到原始DataFrame
        output_df = input_df_original.copy()
        output_df['Predicted_Resistant_Class'] = predicted_classes
        output_df['Predicted_Resistant_Probability'] = predicted_probabilities
        
        # 6. 保存结果
        output_dir = os.path.dirname(args.output_csv)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
            print(f"创建输出目录: {output_dir}")
            
        output_df.to_csv(args.output_csv, index=False)
        print(f"预测结果已保存到: {args.output_csv}")
        print("\n预测结果预览 (前5行):")
        print(output_df.head())

    except FileNotFoundError as e:
        print(f"文件错误: {e}")
    except ValueError as e:
        print(f"值错误: {e}")
    except IOError as e:
        print(f"IO错误: {e}")
    except RuntimeError as e:
        print(f"运行时错误: {e}")
    except Exception as e:
        print(f"预测过程中发生未预料的错误: {e}")
        import traceback
        traceback.print_exc()
